import re
from typing import List

from more_itertools import chunked
from pydantic import BaseModel, field_validator
from tqdm import tqdm
from transformers import AutoTokenizer
from vllm import LLM, SamplingParams


class QA(BaseModel):
    question: str
    answer: str

    @field_validator("question")
    def validate_question(cls, value):
        return value.replace("Q: ", "")

    @field_validator("answer")
    def validate_answer(cls, value):
        return value.replace("A: ", "")


def parse_qas(text: str) -> List[QA]:
    qas = []

    pairs = re.split(r"\n{2,}", text.strip())
    for pair in pairs:
        lines = pair.split("\n")
        question = ""
        answer = ""

        flag = 0
        for line in lines:
            if line.startswith("Q: "):
                flag += 1
            elif line.startswith("A: "):
                flag += 1

            if flag == 1:
                question += line + "\n"
            elif flag == 2:
                answer += line + "\n"

        if flag == 2:
            qas.append(QA(question=question.strip(), answer=answer.strip()))

    return qas


MAX_LEN = 1024


def split_docs(text: str) -> List[str]:
    segments = re.split(r"\n{2,}", text.strip())

    docs = []

    for segment in segments:
        lines = segment.split("\n")
        doc = ""
        for line in lines:
            if len(doc) + len(line) + 1 < MAX_LEN:
                doc += line + "\n"
            else:
                docs.append(doc)
                doc = line + "\n"
        docs.append(doc)

    return docs


# Common prefix.
prefix = """
你需要从慈善组织公开募捐管理办法知识文本中提取QA问答对。
提出的问题和回答应该是有意义的。最好答案可以直接从文本中找到。
问题需要简短一些，针对性强一些。

输出格式：
Q: xxx
A: xxx

Q: xxx
A: xxx

....
"""

example_input = """
第六十九条  建设单位、设计单位、施工单位违反本条例规定，有下列行为之一的，由供水行政主管部门责令改正，并处以五万元以上十万元以下罚款；情节严重的，处以十万元以上三十万元以下罚款：
（一）违反本条例第十条规定，未按照要求设置集中转压、供水管道、消火栓、二次供水等设施，或者供水设施未与建筑物主体工程同时设计、同时施工、同时验收交付使用的；
（二）违反本条例第十一条第一款规定，未按照“一户一表、水表出户”的要求设计、建设住宅项目配建的供水设施的；
（三）违反本条例第十八条第二款规定，未与供水单位商定、实施保护措施，造成公共供水设施损坏的；
（四）违反本条例第十九条第二款、第三款规定，未征求供水单位的意见，未与供水单位商定保护方案，或者改装、迁移后的工程状况和质量低于原工程的。
第七十条  违反本条例第十三条第一款、第十四条第一款、第十五条规定，供水设施建设、竣工验收不符合相关规定的，由供水、住房和城乡建设等行政管理部门依照《建设工程质量管理条例》《建设工程勘察设计管理条例》的有关规定予以处罚。
第七十一条  违反本条例规定，有下列行为之一的，由供水行政主管部门责令改正，可以对个人处以三千元以上三万元以下罚款，对单位处以五千元以上五万元以下罚款；构成治安违法行为的，由公安机关依照《中华人民共和国治安管理处罚法》的规定处理；构成犯罪的，依法追究刑事责任：
（一）违反本条例第十七条规定，在供水设施安全保护范围内，从事禁止性行为的；
（二）违反本条例第十九条第一款规定，擅自拆除、改装、迁移公共供水设施的；
（三）违反本条例第二十二条第一款、第二款规定，未按照规定管理维护供水设施，或者未按二次供水设施清洗保洁技术规范开展二次供水设施清洗、消毒的；
（四）违反本条例第二十三条第三款规定，妨碍供水单位抢修供水设施的；
（五）违反本条例第三十七条第二款、第三款规定，乱收水费的；
（六）违反本条例第四十三条规定，危害供水安全的。
违反本条例第十九条第一款、第四十三条第六项、第七项规定，情节严重的，依照《城市供水条例》的有关规定，经区人民政府批准，可以在一定时间内停止供水。
第七十二条  违反本条例第二十二条第三款、第四款规定，从事二次供水设施清洗、消毒的工作人员不符合从业要求的，由卫生健康行政管理部门责令改正，予以警告；情节严重的，对单位处以五百元以上五千元以下罚款。
第七十三条  违反本条例第四十四条第一款规定的，由供水行政主管部门责令其补缴水费，并处应补缴水费三倍以上五倍以下罚款，最高不超过十万元；违反本条例第四十四条第一款第一项、第四项、第五项、第六项、第七项规定，情节严重的，依照《城市供水条例》的有关规定，经区人民政府批准，可以在一定时间内停止供水。构成治安违法行为的，由公安机关依照《中华人民共和国治安管理处罚法》的规定处理；构成犯罪的，依法追究刑事责任。
"""

example_output = """
Q: 建设单位、设计单位、施工单位违反规定未按要求设置供水设施或未与主体工程同时设计、施工、验收的,会受到什么处罚?
A: 由供水行政主管部门责令改正,并处以5万元以上10万元以下罚款;情节严重的,处以10万元以上30万元以下罚款。

Q: 未按"一户一表、水表出户"要求设计、建设住宅项目配建供水设施的,会受到什么处罚?
A: 由供水行政主管部门责令改正,并处以5万元以上10万元以下罚款;情节严重的,处以10万元以上30万元以下罚款。

Q: 在供水设施安全保护范围内从事禁止性行为的,会受到什么处罚?
A: 由供水行政主管部门责令改正,可以对个人处以3000元以上3万元以下罚款,对单位处以5000元以上5万元以下罚款。情节严重可能构成治安违法或刑事犯罪。

Q: 擅自拆除、改装、迁移公共供水设施的,会受到什么处罚?
A: 由供水行政主管部门责令改正,可以对个人处以3000元以上3万元以下罚款,对单位处以5000元以上5万元以下罚款。情节严重可能构成治安违法或刑事犯罪,甚至可能在一定时间内停止供水。

Q: 违反规定盗用或者擅自用水的,会受到什么处罚?
A: 由供水行政主管部门责令其补缴水费,并处应补缴水费3倍以上5倍以下罚款,最高不超过10万元。情节严重的可能在一定时间内停止供水,构成治安违法或刑事犯罪的将依法追究相应责任。
"""

document = open(".vscode/knowledge.txt", "r", encoding="utf-8").read()
segments = split_docs(document)


# tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-57B-A14B-Instruct")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-7B-Instruct")

generating_prompts = []

for segment in segments:
    messages = [
        {"role": "system", "content": prefix},
        {"role": "user", "content": example_input},
        {"role": "system", "content": example_output},
        {"role": "user", "content": segment},
    ]
    text = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    generating_prompts.append(text)

# Create a sampling params object.
sampling_params = SamplingParams(temperature=0.0, max_tokens=2048)

# Increase gpu_memory_utilization to 0.9
prefix_cached_llm = LLM(
    model="Qwen/Qwen2-7B-Instruct",
    enable_prefix_caching=True,
    gpu_memory_utilization=0.95,
    max_model_len=8192,
)

# Warmup so that the shared prompt's KV cache is computed.
prefix_cached_llm.generate(generating_prompts[0], sampling_params, use_tqdm=False)

# Generate with prefix caching.
outputs = []
for batch in tqdm(chunked(generating_prompts, 8), total=len(generating_prompts) // 8):
    outputs += prefix_cached_llm.generate(batch, sampling_params, use_tqdm=False)

results = []
# Print the outputs. You should see the same outputs as before.
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    qas = parse_qas(generated_text)
    results.extend(qas)


with open(".vscode/faq.txt", "w", encoding="utf-8") as f:
    for qa in results:
        f.write(f"Q: {qa.question}\nA: {qa.answer}\n\n")
